/*
 * Copyright 2010-2012 Susanta Tewari. <freecode4susant@users.sourceforge.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package coalescent.is;

import bd.commons.math.distribution.SampleSpace;
import coalescent.Event;
import coalescent.EventType;
import coalescent.Genealogy;
import coalescent.model.PopGenModel;
import coalescent.statistic.AC;
import commons.is.Proposal_w_Prob;

import java.lang.reflect.Array;
import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.logging.Logger;

import static commons.util.ExceptionUtil.throwArgEx;

/**
 * Proposal for genealogy built in a sequential manner from the sample to the MRCA. Each step in
 * the sequence corresponds to an event in the genealogy.
 *
 * @param <S> sample type
 * @param <M> model type
 * @author Susanta Tewari
 * @version version
 * @history First created on 06/02/2011.
 * @since since
 */
public abstract class GProposal<S extends AC<S, M>, M extends PopGenModel>
        implements Proposal_w_Prob<Genealogy<S, M>> {

    protected final S sample;
    private Map<S, ElementSampler> cache;
    private static final Logger LOGGER = Logger.getLogger(GProposal.class.getName());

    /**
     * @param sample sample
     */
    public GProposal(final S sample) {

        this.sample = sample;
        this.cache  = new HashMap<>(sample.eventsToMRCA());
    }

    /**
     * Specifies the proposal weight for the given configuration, allele and event type.
     *
     * @param config sample configuration
     * @param allele allele on the current configuration
     * @param eventType the chosen event type
     * @return the corresponding proposal probability
     */
    protected abstract BigDecimal proposalWeight(S config, Object allele, EventType eventType);

    @Override
    public final Genealogy<S, M> sample() {

        cache = new HashMap<>(sample.eventsToMRCA());

        final List<Event<S>> eventChain = new ArrayList<>(sample.eventsToMRCA());
        S config                        = sample;

        while (!config.isMRCA()) {

            final ElementSampler elementSampler = createElementSampler(config);

            cache.put(config, elementSampler);

            final Element e         = elementSampler.propose();
            final S ancestralConfig = config.apply(e.getAllele(), e.getEventType());

            eventChain.add(new Event<>(config, ancestralConfig, e.getAllele(), e.getEventType()));

            config = ancestralConfig;
        }

        return new Genealogy<>(eventChain);
    }

    @Override
    public final BigDecimal probability(final Genealogy<S, M> x) {

        final List<Event<S>> event_chain = x.getEventChain();
        BigDecimal result                = BigDecimal.ONE;

        for (final Event<S> event : event_chain) {

            final Object allele                 = event.getPreAllele();
            final EventType type                = event.getType();
            final S config                      = event.getPre();
            final ElementSampler elementSampler = cache.get(config);
            final BigDecimal proposal_prob      = elementSampler.getProbability(allele, type);

            result = result.multiply(proposal_prob);
        }

        return result;
    }

    private ElementSampler createElementSampler(final S config) {

        final List<Element> e_list = new ArrayList<>(config.getN());

        for (final EventType type : getModel().getEventTypes()) {

            for (final Object allele : config.alleles(type)) {
                e_list.add(new Element(config, allele, type));
            }
        }

        if (e_list.isEmpty()) {

            LOGGER.severe("config w/ missing alleles. config: " + config);

            throw new IllegalStateException("config w/ missing alleles");
        }

        final Element[] e_array = e_list.toArray((Element[]) Array.newInstance(Element.class, 0));

        return new ElementSampler(e_array);
    }

    protected final M getModel() {
        return sample.getModel();
    }

    /**
     * An element of the proposal. It contains enough information to generate its ancestral and its
     * proposal weight.
     */
    private final class Element {

        private final Object allele;
        private final EventType eventType;
        private final S config;

        protected Element(final S config, final Object allele, final EventType eventType) {

            this.config    = config;
            this.allele    = allele;
            this.eventType = eventType;
        }

        public Object getAllele() {
            return allele;
        }

        public EventType getEventType() {
            return eventType;
        }

        public BigDecimal getWeight() {
            return GProposal.this.proposalWeight(config, allele, eventType);
        }
    }

    /**
     * Generates a random sample from a set of elements with weights.
     * todo scope of optimization: store proposed value and its prob.
     */
    private final class ElementSampler {

        private final Element[] elements;

        ElementSampler(final Element[] elements) {

            throwArgEx(elements.length == 0, "empty elements");

            this.elements = elements;
        }

        private SampleSpace<Element> createSampleSpace() {

            final BigDecimal[] proposalWeights = new BigDecimal[elements.length];

            for (int i = 0; i < proposalWeights.length; i++) {
                proposalWeights[i] = elements[i].getWeight();
            }

            return new SampleSpace<>(elements, proposalWeights);
        }

        Element propose() {
            return createSampleSpace().sample();
        }

        BigDecimal getProbability(final Object allele, final EventType eventType) {

            for (final Element element : elements) {

                if (element.getAllele().equals(allele) && (element.getEventType() == eventType)) {
                    return createSampleSpace().elementProbability(element);
                }
            }

            throw new IllegalArgumentException(
                "no proposal found with the specified allele and event type");
        }
    }
}
